package arithmetic2.day10;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;

/*
    并查集结构
 */
public class UnionFind<T> {
    //存储每个元素和它对应的节点关系
    private Map<T, Node<T>> nodeMap = new HashMap<>();
    //存储一个节点和它父节点的对应关系
    private Map<Node<T>, Node<T>> fatherMap = new HashMap<>();
    //存储代表和它所在集合的大小关系
    private Map<Node<T>, Integer> sizeMap = new HashMap<>();

    public UnionFind(List<T> list) {
        //将所有元素包装为Node类型,添加到三个map中
        for (T t : list) {
            Node<T> node = new Node<>(t);
            nodeMap.put(t, node);
            fatherMap.put(node, node);
            sizeMap.put(node, 1);
        }
    }

    /*
        找到一个节点的代表节点：
            准备一个栈，将沿途所有的节点都直接放在代表节点下,优化并查集结构
     */
    public Node<T> findHead(Node<T> node) {
        Stack<Node<T>> stack = new Stack<>();
        while (fatherMap.get(node) != node) {
            //将这个父节压栈
            stack.push(fatherMap.get(node));
            node = fatherMap.get(node);
        }
        //将栈中所有元素直接放在代表节点下
        while (!stack.isEmpty()) {
            fatherMap.put(stack.pop(), node);
        }
        return node;
    }

    /*
        查询两个节点是否属于同一集合
     */
    public boolean isSameSet(T t1, T t2) {
        //两个节点都存在
        if (nodeMap.containsKey(t1) && nodeMap.containsKey(t2)) {
            return findHead(nodeMap.get(t1)) == findHead(nodeMap.get(t2));
        }
        return false;
    }

    /*
        将两个节点所在的集合合并
     */
    public void union(T t1, T t2) {
        //两个元素都存在并且不在同一集合才需要合并
        if (nodeMap.containsKey(t1) && nodeMap.containsKey(t2) && !isSameSet(t1, t2)) {
            //找到两个节点的代表节点
            Node<T> t1Head = findHead(nodeMap.get(t1));
            Node<T> t2Head = findHead(nodeMap.get(t2));
            //集合大的代表节点
            Node<T> big = sizeMap.get(t1Head) > sizeMap.get(t2Head) ? t1Head : t1Head;
            //集合小的代表节点
            Node<T> small = big == t1Head ? t2Head : t1Head;
            //将小的代表节点放在大的节点下
            fatherMap.put(small, big);
            //设置集合大小
            sizeMap.put(big, sizeMap.get(big) + sizeMap.get(small));
            //将较小的代表节点再sizeMap中删除
            sizeMap.remove(small);
        }
    }

    /*
         并查集的节点类型
     */
    public static class Node<T> {
        public T value;

        public Node(T value) {
            this.value = value;
        }
    }
}
